import json
import os
import re
import pandas as pd
from tqdm import tqdm

path_dir = os.path.dirname(os.path.abspath(__file__))

def build_prop_to_ref(props_with_ref):
    prop2ref = dict()
    for prop_with_ref in props_with_ref:
        if prop_with_ref['define_with_ref']:
            prop2ref[prop_with_ref['prop']] = {
                'define_with_ref': prop_with_ref['define_with_ref'],
                'ref_url': prop_with_ref['ref']['url'],
                'ref_content': prop_with_ref['ref']['content'],
            }
    return prop2ref

if __name__ == '__main__':
    phase25_model = 'qwen3-max'
    with open(os.path.join(path_dir, "phase2.5_output", f"{phase25_model}_prop_desc.jsonl"), 'r', encoding='utf-8') as fp:
        lines = fp.readlines()
        props_with_ref = [json.loads(line) for line in lines]
    prop2ref = build_prop_to_ref(props_with_ref)

    with open(os.path.join(path_dir, "phase6_output", "qwen3_reranker_verify_expand_questions.jsonl"), 'r', encoding='utf-8') as fp:
        lines = fp.readlines()
        ex_questions_data = [json.loads(line) for line in lines]

    threshold = 0.85
    filtered_questions_data = []
    for ex_question_data in tqdm(ex_questions_data):
        ex_question_data.pop('meta')
        # 1. 判断原始question是否合格
        if ex_question_data['score'] < threshold:
            continue

        # 2. 判断扩展question是否合格，不合格会过滤掉
        new_questions = ex_question_data.pop('new_questions')
        new_questions_score = ex_question_data.pop('new_questions_score')
        assert len(new_questions) == len(new_questions_score)
        questions_cnt = len(new_questions)
        filtered_questions_data_and_score = [(c, s) for c, s in zip(new_questions, new_questions_score) if s>=threshold]
        for i in range(questions_cnt):
            if i < len(filtered_questions_data_and_score):
                c, s = filtered_questions_data_and_score[i]
                ex_question_data[f'expand_question_{i}'] = c
                ex_question_data[f'expand_question_score_{i}'] = round(s, 2)
            else:
                ex_question_data[f'expand_question_{i}'] = None
                ex_question_data[f'expand_question_score_{i}'] = None

        # 3. 判断是否是定义题，定义题给出原始的ref相关信息
        if ex_question_data['rule'] == 'RULE-A':
            ref = prop2ref[ex_question_data['prop']]
            ex_question_data['define_with_ref'] = ref['define_with_ref']
            ex_question_data['ref_url'] = ref['ref_url']
            ex_question_data['ref_content'] = ref['ref_content']
        else:
            ex_question_data['define_with_ref'] = None
            ex_question_data['ref_url'] = None
            ex_question_data['ref_content'] = None

        filtered_questions_data.append(ex_question_data)

    with open(os.path.join(path_dir, "phase7_output", f"Expand_QA_threshold_{threshold}.json"), 'w', encoding='utf-8') as fp:
        json.dump(filtered_questions_data, fp, ensure_ascii=False, indent=4)

    df = pd.DataFrame(data=filtered_questions_data)
    df.to_excel(os.path.join(path_dir, "phase7_output", f"Expand_QA_threshold_{threshold}.xlsx"), index=False)
